package com.lagou.utils;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @author Roger
 * @date 2020/8/15 9:49
 * @description
 */
@Component
public class ConnectionUtil {

    private ThreadLocal<Connection> threadLocal = new ThreadLocal<>();

    @Autowired
    private DataSource dataSource;

    public Connection getConnection() {
        Connection connection = null;
        try {
            if (threadLocal.get() == null) {
                connection = dataSource.getConnection();
                threadLocal.set(connection);
            } else {
                connection = threadLocal.get();
            }
        } catch (SQLException e) {
            e.printStackTrace();
        }
        return connection;
    }

    public void removeThreadConnection() {
        threadLocal.remove();
    }
}
